import argparse
import base64
import json
import os
import os.path as osp

import imgviz
import PIL.Image

from labelme import utils
from labelme.logger import logger


def main():
    logger.warning(
        "DEPRECATED: This script will be removed in the near future. "
        "Please use `labelme_export_json` instead."
    )
    logger.warning(
        "NOTE: This script is aimed to demonstrate how to convert a JSON file "
        "to a single image dataset. so it won't handle multiple JSON files to "
        "generate a real-use dataset."
    )

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", default="input")
    parser.add_argument("--output_dir", default="output")
    args = parser.parse_args()

    input_dir = args.input_dir
    output_dir = args.output_dir

    if not osp.exists(output_dir):
        os.mkdir(output_dir)

    for json_file in os.listdir(input_dir):
        if json_file.endswith(".json"):
            data = json.load(open(osp.join(input_dir, json_file)))
            imageData = data.get("imageData")

            if not imageData:
                imagePath = os.path.join(os.path.dirname(json_file), data["imagePath"])
                with open(imagePath, "rb") as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode("utf-8")
            img = utils.img_b64_to_arr(imageData)

            label_name_to_value = {"_background_": 0}
            for shape in sorted(data["shapes"], key=lambda x: x["label"]):
                label_name = shape["label"]
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value
            lbl, _ = utils.shapes_to_label(img.shape, data["shapes"], label_name_to_value)

            label_names = [None] * (max(label_name_to_value.values()) + 1)
            for name, value in label_name_to_value.items():
                label_names[value] = name

            lbl_viz = imgviz.label2rgb(
                lbl, imgviz.asgray(img), label_names=label_names, loc="rb"
            )

            output_subdir = osp.join(output_dir, osp.splitext(json_file)[0])
            if not osp.exists(output_subdir):
                os.mkdir(output_subdir)

            PIL.Image.fromarray(img).save(osp.join(output_subdir, "img.png"))
            utils.lblsave(osp.join(output_subdir, "label.png"), lbl)
            PIL.Image.fromarray(lbl_viz).save(osp.join(output_subdir, "label_viz.png"))

            with open(osp.join(output_subdir, "label_names.txt"), "w") as f:
                for lbl_name in label_names:
                    f.write(lbl_name + "")

            logger.info("Saved to: {}".format(output_subdir))


if __name__ == "__main__":
    main()
